
#include "StdAfx.h"
#include "DetourManager.h"
#include "..\Detours\detours.h"


#include <psapi.h>


CStreamFormatter CDetourManager::m_log;
CWriteStreamOnFile CDetourManager::m_ws;

#pragma data_seg("Shared")
HHOOK CDetourManager::g_hHook = NULL;
HINSTANCE   CDetourManager::g_hInstance = NULL;         // This instance
#pragma data_seg()

// Function pointers...

static int (WINAPI * TrueGetWindowTextA)(HWND hWnd, LPSTR lpString, int nMaxCount) = GetWindowTextA;
static int (WINAPI * TrueGetWindowTextW)(HWND hWnd, LPWSTR lpString, int nMaxCount) = GetWindowTextW;
static HANDLE (WINAPI * TrueOpenProcess)(DWORD, BOOL, DWORD) = OpenProcess;
static BOOL (WINAPI * TrueCreateProcessW)(LPCWSTR, LPWSTR, LPSECURITY_ATTRIBUTES, LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCWSTR, LPSTARTUPINFOW, LPPROCESS_INFORMATION) = CreateProcessW;
static BOOL (WINAPI * TrueCreateProcessA)(LPCSTR, LPSTR, LPSECURITY_ATTRIBUTES, LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCSTR, LPSTARTUPINFOA, LPPROCESS_INFORMATION) = CreateProcessA;
static BOOL (WINAPI * TrueEnumProcesses)(DWORD*, DWORD, DWORD*) = EnumProcesses;
static BOOL (WINAPI * TrueReadProcessMemory)(HANDLE, LPCVOID, LPVOID, SIZE_T, SIZE_T*) = ReadProcessMemory;
static HANDLE (WINAPI * TrueCreateFileW)(LPCWSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile) = CreateFileW;
static HANDLE (WINAPI * TrueCreateFileA)(LPCSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile) = CreateFileA;

CDetourManager::CDetourManager(void)
{
}

CDetourManager::~CDetourManager(void)
{
}

int WINAPI MyGetWindowTextA(HWND hWnd, LPSTR lpString, int nMaxCount)
{
	int retVal = TrueGetWindowTextA(hWnd, lpString, nMaxCount);
	if (retVal != 0)
	{
		USES_CONVERSION;
		LPCTSTR lpts = A2T(lpString);
		CDetourManager::Send(_T("GetWindowTextA(hWnd=%p,lpString='%s',nMaxCount=%d)"), hWnd, lpts, nMaxCount);
	}
	return retVal;

}

int WINAPI MyGetWindowTextW(HWND hWnd, LPWSTR lpString, int nMaxCount)
{
	int retVal = TrueGetWindowTextW(hWnd, lpString, nMaxCount);
	if (retVal != 0)
	{
		//ApplicationProxy::Send(lpString, hWnd);
	}

	return retVal;

}

HANDLE WINAPI MyOpenProcess(DWORD dwDesiredAccess, BOOL bInheritHandle, DWORD dwProcessID)
{
	HANDLE retVal = TrueOpenProcess(dwDesiredAccess, bInheritHandle, dwProcessID);
	
	return retVal;

}

BOOL WINAPI MyCreateProcessW(LPCWSTR lpAppName, LPWSTR lpCommandLine, LPSECURITY_ATTRIBUTES lpProcessAttributes,
							 LPSECURITY_ATTRIBUTES lpThreadAttributes, BOOL bInheritHandles, DWORD dwCreationFlags,
							 LPVOID lpEnvironment, LPCWSTR lpCurrentDirectory, LPSTARTUPINFOW lpStartupInfo,
							 LPPROCESS_INFORMATION lpProcessInfo)
{
	BOOL bRet = TrueCreateProcessW(lpAppName, lpCommandLine, lpProcessAttributes,
									lpThreadAttributes, bInheritHandles, dwCreationFlags,
									lpEnvironment, lpCurrentDirectory, lpStartupInfo,
									lpProcessInfo);

	return bRet;

}

BOOL WINAPI MyCreateProcessA(LPCSTR lpAppName, LPSTR lpCommandLine, LPSECURITY_ATTRIBUTES lpProcessAttributes,
							 LPSECURITY_ATTRIBUTES lpThreadAttributes, BOOL bInheritHandles, DWORD dwCreationFlags,
							 LPVOID lpEnvironment, LPCSTR lpCurrentDirectory, LPSTARTUPINFOA lpStartupInfo,
							 LPPROCESS_INFORMATION lpProcessInfo)
{
	BOOL bRet = TrueCreateProcessA(lpAppName, lpCommandLine, lpProcessAttributes,
									lpThreadAttributes, bInheritHandles, dwCreationFlags,
									lpEnvironment, lpCurrentDirectory, lpStartupInfo,
									lpProcessInfo);

	return bRet;
}

BOOL WINAPI MyEnumProcesses(DWORD* pProcessIDs, DWORD cb, DWORD* bytesReturned)
{
	BOOL bRet = TrueEnumProcesses(pProcessIDs, cb, bytesReturned);

	return bRet;
}

BOOL WINAPI MyReadProcessMemory(HANDLE hProcess, LPCVOID lpBaseAddress, LPVOID lpBuffer, SIZE_T nSize, SIZE_T* lpNumberOfBytesRead)
{
	BOOL bRet = TrueReadProcessMemory(hProcess, lpBaseAddress, lpBuffer, nSize, lpNumberOfBytesRead);

	return bRet;
}

HANDLE MyCreateFileW(LPCWSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile)
{
	HANDLE hRet = TrueCreateFileW(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);

	return hRet;
}

HANDLE MyCreateFileA(LPCSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile)
{
	HANDLE hRet = TrueCreateFileA(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);

	return hRet;
}


HRESULT CDetourManager::Init(HANDLE hModule)
{
	HRESULT hr;
	USES_CONVERSION;

	g_hInstance = (HINSTANCE)hModule;

	if(FAILED(hr=m_ws.Initialize(_T("Dtour.log"), true)))
		return hr;

	m_log.Initialize(&m_ws);
	m_log.WriteFormatted(_T("Watch.dll log opened "));

	COleDateTime now = COleDateTime::GetCurrentTime();
	m_log.WriteRaw(now.Format(_T("%H:%M:%S %A, %B %d, %Y")));
	m_log.WriteRaw(_T("\r\n"));

	InstallDetour((void**)&TrueGetWindowTextA, MyGetWindowTextA);
	InstallDetour((void**)&TrueGetWindowTextW, MyGetWindowTextW);
	//InstallDetour((void**)&TrueOpenProcess, MyOpenProcess);
	//InstallDetour((void**)&TrueCreateProcessW, MyCreateProcessW);
	//InstallDetour((void**)&TrueCreateProcessA, MyCreateProcessA);
	//InstallDetour((void**)&TrueEnumProcesses, MyEnumProcesses);
	//InstallDetour((void**)&TrueReadProcessMemory, MyReadProcessMemory);
	//InstallDetour((void**)&TrueCreateFileW, MyCreateFileW);
	//InstallDetour((void**)&TrueCreateFileA, MyCreateFileA);

	g_hHook = SetWindowsHookEx(WH_CBT, (HOOKPROC) CBTProc, g_hInstance, 0);

	m_log.WriteFormatted(_T("Hook %p installed to entry=%p, hInstance=%p\r\n"), g_hHook, CBTProc, g_hInstance);

	m_ws.FlushStream();

	return S_OK;
}


void CDetourManager::Exit()
{
	UninstallDetour(TrueGetWindowTextA, MyGetWindowTextA);
	UninstallDetour(TrueGetWindowTextW, MyGetWindowTextW);
	//UninstallDetour(TrueOpenProcess, MyOpenProcess);
	//UninstallDetour(TrueCreateProcessW, MyCreateProcessW);
	//UninstallDetour(TrueCreateProcessA, MyCreateProcessA);
	//UninstallDetour(TrueEnumProcesses, MyEnumProcesses);
	//UninstallDetour(TrueReadProcessMemory, MyReadProcessMemory);
	//UninstallDetour(TrueCreateFileW, MyCreateFileW);
	//UninstallDetour(TrueCreateFileA, MyCreateFileA);

	if(NULL!=g_hHook)
	{
		m_log.WriteFormatted(_T("Hook %p uninstalled for entry %p, hInstance=%p\r\n",g_hHook,CBTProc,g_hInstance ));
		UnhookWindowsHookEx(g_hHook);
		g_hHook = NULL;
	}
}


bool CDetourManager::InstallDetour(void** originalFunction, void* newFunction)
{
	LONG lError = DetourTransactionBegin();
	if (lError != NO_ERROR)
		return false;

    lError = DetourUpdateThread(GetCurrentThread());
	if (lError != NO_ERROR)
	{
		DetourTransactionAbort();
		return false;
	}

    lError = DetourAttach(originalFunction, newFunction);
	if (lError != NO_ERROR)
	{
		DetourTransactionAbort();
		return false;
	}

    lError = DetourTransactionCommit();
	if (lError != NO_ERROR)
	{
		DetourTransactionAbort();
		return false;
	}

	return true;
}


bool CDetourManager::UninstallDetour(void* originalFunction, void* newFunction)
{
    DetourTransactionBegin();
    DetourUpdateThread(GetCurrentThread());
    DetourDetach(&originalFunction, newFunction);
    DetourTransactionCommit();

	return true;
}



void CDetourManager::Send(LPCTSTR lpFormat,...)
{
	va_list ap;
	va_start(ap,lpFormat);

	CString str;

	str.FormatV(lpFormat, ap);
	m_log.WriteRaw(str);
	m_log.WriteRaw(_T("\r\n"));

	va_end(ap);

}


///////////////////////////////////////////////////////////////////////////////
// This is the CBT hook procedure.  We use a Windows-based CBT hook not only
// to inject our DLL into the process, but to detect when top-level windows -
// such as poker table windows - are created.
///////////////////////////////////////////////////////////////////////////////
LRESULT CALLBACK CDetourManager::CBTProc(int nCode, WPARAM wParam, LPARAM lParam)
{
	if (nCode < 0)
		return CallNextHookEx(g_hHook, nCode, wParam, lParam);
	

	// Window was activated. We currently use this to detect window creation also.
	// Why not use HCBT_CREATE? Because often the window doesn't yet have a caption
	// when HCBT_CREATE is called, and our sample window-detection logic looks
	// at the window caption to determine the type.
	HWND hWnd = (HWND)wParam;
	if (!hWnd)
		return 0;

	if (nCode == HCBT_ACTIVATE)
	{
		//if (!g_pClient->IsRegisteredWindow(hWnd))
		//    g_pClient->TryRegisterWindow(hWnd);
		CDetourManager::Send(_T("HCBT_ACTIVATE(hWnd=%p,lParam=%lx"),hWnd,lParam);
	}
	else if (nCode == HCBT_DESTROYWND)
	{
		//if (g_pClient->IsRegisteredWindow(hWnd))
		//    g_pClient->UnregisterWindow(hWnd);
		CDetourManager::Send(_T("HCBT_DESTROYWND(hWnd=%p,lParam=%lx"),hWnd,lParam);
	}else
	{
		CDetourManager::Send(_T("CBT(%p) wParam=%p lParam=%lx"), wParam, lParam);
	}

	// Return 0 to allow window creation/destruction/activation to proceed as normal.
	return 0;
}